#!/usr/bin/env python3
# -*- coding: utf8 -*-

import sys
import random

from Bio import SeqIO, SeqRecord

import config

SEGMENTS_COUNT = config.SEGMENTS_COUNT
CHROMOSOME_FILENAME = config.CHROMOSOME_FILENAME
STRUCTURE_FILENAME = config.STRUCTURE_FILENAME
SCAFFOLDS_FILENAME = config.SCAFFOLDS_FILENAME


def main():
    chromosome = load_chromosome()
    separators = get_separators_list(len(chromosome.seq))
    chromosome_segments = make_chromosome_segments(chromosome, separators)
    save_segments(chromosome_segments)


def load_chromosome():
    sys.stderr.write('Load chromosome...\n')
    with open(CHROMOSOME_FILENAME) as handle:
        chromosome = SeqIO.read(handle, 'fasta')
    sys.stderr.write('Done.\n')
    return chromosome


def get_separators_list(chromosome_length):
    spaces = calc_spaces(chromosome_length)
    sys.stderr.write('Set separators\n')
    separators = []
    for i in range(SEGMENTS_COUNT - 1):
        segment = random.choice(spaces)
        separators.append(random.randint(segment[0], segment[1]))
    separators = sorted(separators)
    return separators


def calc_spaces(chromosome_length):
    exon_segments = load_exon_segments_db()
    segments = filter_exon_segments(exon_segments)
    sys.stderr.write('Calc spaces\n')
    spaces = []
    start = 0
    for segment in segments:
        spaces.append((start, segment[0]))
        start = segment[1]
    spaces.append((start, chromosome_length))
    return spaces


def filter_exon_segments(exon_segments):
    sys.stderr.write('Calc segments\n')
    segments = []
    it = iter(exon_segments)
    prev = next(it)
    start = prev[0]
    for exon in it:
        if prev[1] < exon[0]:
            segments.append((start, prev[1]))
            start = exon[0]
        prev = exon
    return segments


def load_exon_segments_db():
    """ sorted [(exon_start, exon_end)] """
    sys.stderr.write('Load exon database...\n')
    exon_segments = []
    with open(STRUCTURE_FILENAME) as handle:
        handle.readline()  # skip header line
        for line in handle.readlines():
            e_start, e_end = map(int, line.split()[2:4])
            exon_segments.append((e_start - 1, e_end))
    exon_segments = sorted(exon_segments)
    sys.stderr.write('Done.\n')
    return exon_segments


def make_chromosome_segments(chr15, separators):
    sys.stderr.write('Create chromosome segments...\n')
    results = []
    start = 0
    for separator in separators:
        results.append(chr15.seq[start:separator])
        start = separator
    results.append(chr15.seq[start:len(chr15.seq)])
    sys.stderr.write('Done.\n')
    return results


def save_segments(results):
    sys.stderr.write('Save to file\n')
    with open(SCAFFOLDS_FILENAME, 'w') as handle:
        sequences = (SeqRecord.SeqRecord(seq, id='segm' + str(i), description='') for i, seq in enumerate(results))
        SeqIO.write(sequences, handle, 'fasta')
    sys.stderr.write('Done.\n')


if __name__ == '__main__':
    main()
